import base64
import secrets
from typing import TypeAlias

from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
from cryptography.hazmat.primitives.kdf.hkdf import HKDF

# 类型别名
StrBase64: TypeAlias = str

# 常量
NONCE_SIZE = 12
KEY_SIZE = 32
DEFAULT_MASTER_KEY = b"hello"
DEFAULT_HKDF_INFO = b"session_key_derivation"


class Chacha20CryptoError(Exception):
    """加密/解密操作相关的异常"""

    pass


def chacha20_encrypt_b64(plain_text: str, key: str) -> StrBase64:
    k = key.encode()
    if len(k) != KEY_SIZE:
        raise ValueError("Key must be exactly 32 bytes long. current:" + str(len(k)))
    data = _chacha20_encrypt_(plain_text.encode(), k)
    return base64.b64encode(data).decode()


def chacha20_decrypt_b64(encrypted_b64_text: StrBase64, key: str) -> str:
    k = key.encode()
    if len(k) != KEY_SIZE:
        raise ValueError("Key must be exactly 32 bytes long. current:" + str(len(k)))
    data = _chacha20_decrypt_(base64.b64decode(encrypted_b64_text), k)
    return data.decode()


def generate_hkdf_key(
    dynamic_key: str, master_key: bytes = DEFAULT_MASTER_KEY, length: int = KEY_SIZE
) -> str:
    """内部计算函数"""
    hkdf = HKDF(
        algorithm=hashes.SHA256(),
        length=KEY_SIZE,
        salt=None,
        info=DEFAULT_HKDF_INFO,
    )
    k = hkdf.derive(master_key + dynamic_key.encode("utf-8"))
    text = base64.b64encode(k).decode("utf-8")
    return text[:length]


def _chacha20_encrypt_(data: bytes, key: bytes) -> bytes:
    """内部加密函数"""
    chacha = ChaCha20Poly1305(key)
    nonce = secrets.token_bytes(NONCE_SIZE)
    ciphertext = chacha.encrypt(nonce, data, None)
    return nonce + ciphertext


def _chacha20_decrypt_(encrypted_data: bytes, key: bytes) -> bytes:
    """内部解密函数"""
    chacha = ChaCha20Poly1305(key)
    nonce, ciphertext = encrypted_data[:NONCE_SIZE], encrypted_data[NONCE_SIZE:]
    return chacha.decrypt(nonce, ciphertext, None)


def test() -> None:
    """测试函数"""
    try:
        key = generate_hkdf_key("123456")
        print(f"Generated key: {key}")
        plaintext = "This is a secret server message."
        print(f"Original text: {plaintext}")
        encrypted = chacha20_encrypt_b64(plaintext, key)
        print(f"Encrypted text: {encrypted}")
        decrypted = chacha20_decrypt_b64(encrypted, key)
        print(f"Decrypted text: {decrypted}")
        assert plaintext == decrypted, "Decryption failed: text doesn't match original"
        print("✓ Encryption/decryption test passed!")
    except Chacha20CryptoError as e:
        print(f"Error: {e}")


if __name__ == "__main__":
    test()
